//! 逻辑运算指令编码模块
//! 
//! 实现 8086 逻辑运算指令：AND, OR, XOR, NOT, TEST

use crate::ast::Operand;
use crate::error::{AssemblerError, Result};
use super::registers::RegisterEncoding;

/// 逻辑运算指令编码器
pub struct LogicalEncoder;

impl LogicalEncoder {
    /// 编码逻辑运算指令
    pub fn encode(&self, instruction: &str, operands: &[Operand]) -> Result<Vec<u8>> {
        match instruction {
            "与" => self.encode_and(operands),
            "或" => self.encode_or(operands),
            "异或" => self.encode_xor(operands),
            "非" => self.encode_not(operands),
            "测试" => self.encode_test(operands),
            _ => Err(AssemblerError::CodeGenError {
                message: format!("未知的逻辑指令: {}", instruction),
            }),
        }
    }

    /// 编码与指令（AND）
    fn encode_and(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 2 {
            return Err(AssemblerError::OperandError {
                message: "与 指令需要 2 个操作数".to_string(),
                line: 0,
            });
        }

        match (&operands[0], &operands[1]) {
            (Operand::Register(dst), Operand::Register(src)) => {
                let dst_code = RegisterEncoding::get_16bit_register_code(dst)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", dst),
                        line: 0,
                    })?;
                let src_code = RegisterEncoding::get_16bit_register_code(src)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", src),
                        line: 0,
                    })?;
                let opcode = 0x21; // AND reg, reg (16位)
                let modrm = (0b11 << 6) | (dst_code << 3) | src_code;
                Ok(vec![opcode, modrm])
            }
            (Operand::Register(reg), Operand::Immediate(imm)) => {
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                if *imm >= -128 && *imm <= 127 {
                    // 8位立即数
                    let opcode = 0x83; // AND reg, imm8
                    let modrm = (0b11 << 6) | (0b100 << 3) | reg_code; // reg字段=100表示AND
                    Ok(vec![opcode, modrm, *imm as u8])
                } else {
                    // 16位立即数
                    let opcode = 0x81; // AND reg, imm16
                    let modrm = (0b11 << 6) | (0b100 << 3) | reg_code; // reg字段=100表示AND
                    let imm_u16 = *imm as u16;
                    Ok(vec![opcode, modrm, imm_u16 as u8, (imm_u16 >> 8) as u8])
                }
            }
            _ => Err(AssemblerError::OperandError {
                message: "与暂不支持此操作数组合".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码或指令（OR）
    fn encode_or(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 2 {
            return Err(AssemblerError::OperandError {
                message: "或 指令需要 2 个操作数".to_string(),
                line: 0,
            });
        }

        match (&operands[0], &operands[1]) {
            (Operand::Register(dst), Operand::Register(src)) => {
                let dst_code = RegisterEncoding::get_16bit_register_code(dst)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", dst),
                        line: 0,
                    })?;
                let src_code = RegisterEncoding::get_16bit_register_code(src)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", src),
                        line: 0,
                    })?;
                let opcode = 0x09; // OR reg, reg (16位)
                let modrm = (0b11 << 6) | (dst_code << 3) | src_code;
                Ok(vec![opcode, modrm])
            }
            (Operand::Register(reg), Operand::Immediate(imm)) => {
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                if *imm >= -128 && *imm <= 127 {
                    // 8位立即数
                    let opcode = 0x83; // OR reg, imm8
                    let modrm = (0b11 << 6) | (0b001 << 3) | reg_code; // reg字段=001表示OR
                    Ok(vec![opcode, modrm, *imm as u8])
                } else {
                    // 16位立即数
                    let opcode = 0x81; // OR reg, imm16
                    let modrm = (0b11 << 6) | (0b001 << 3) | reg_code; // reg字段=001表示OR
                    let imm_u16 = *imm as u16;
                    Ok(vec![opcode, modrm, imm_u16 as u8, (imm_u16 >> 8) as u8])
                }
            }
            _ => Err(AssemblerError::OperandError {
                message: "或暂不支持此操作数组合".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码异或指令（XOR）
    fn encode_xor(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 2 {
            return Err(AssemblerError::OperandError {
                message: "异或 指令需要 2 个操作数".to_string(),
                line: 0,
            });
        }

        match (&operands[0], &operands[1]) {
            (Operand::Register(dst), Operand::Register(src)) => {
                let dst_code = RegisterEncoding::get_16bit_register_code(dst)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", dst),
                        line: 0,
                    })?;
                let src_code = RegisterEncoding::get_16bit_register_code(src)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", src),
                        line: 0,
                    })?;
                let opcode = 0x31; // XOR reg, reg (16位)
                let modrm = (0b11 << 6) | (dst_code << 3) | src_code;
                Ok(vec![opcode, modrm])
            }
            (Operand::Register(reg), Operand::Immediate(imm)) => {
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                if *imm >= -128 && *imm <= 127 {
                    // 8位立即数
                    let opcode = 0x83; // XOR reg, imm8
                    let modrm = (0b11 << 6) | (0b110 << 3) | reg_code; // reg字段=110表示XOR
                    Ok(vec![opcode, modrm, *imm as u8])
                } else {
                    // 16位立即数
                    let opcode = 0x81; // XOR reg, imm16
                    let modrm = (0b11 << 6) | (0b110 << 3) | reg_code; // reg字段=110表示XOR
                    let imm_u16 = *imm as u16;
                    Ok(vec![opcode, modrm, imm_u16 as u8, (imm_u16 >> 8) as u8])
                }
            }
            _ => Err(AssemblerError::OperandError {
                message: "异或暂不支持此操作数组合".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码非指令（NOT）
    fn encode_not(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 1 {
            return Err(AssemblerError::OperandError {
                message: "非 指令需要 1 个操作数".to_string(),
                line: 0,
            });
        }

        if let Operand::Register(reg) = &operands[0] {
            let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                .ok_or_else(|| AssemblerError::OperandError {
                    message: format!("无效的16位寄存器: {}", reg),
                    line: 0,
                })?;
            let opcode = 0xF7; // NOT reg (16位)
            let modrm = (0b11 << 6) | (0b010 << 3) | reg_code; // reg字段=010表示NOT
            Ok(vec![opcode, modrm])
        } else {
            Err(AssemblerError::OperandError {
                message: "非只支持寄存器操作数".to_string(),
                line: 0,
            })
        }
    }

    /// 编码测试指令（TEST）
    /// 
    /// TEST 指令与 AND 指令类似，但不修改操作数，只设置标志位
    /// 支持的操作数组合：
    /// - TEST reg, reg
    /// - TEST reg, imm
    /// - TEST reg, mem
    /// - TEST mem, reg
    fn encode_test(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 2 {
            return Err(AssemblerError::OperandError {
                message: "测试 指令需要 2 个操作数".to_string(),
                line: 0,
            });
        }

        match (&operands[0], &operands[1]) {
            // TEST reg, reg
            (Operand::Register(dst), Operand::Register(src)) => {
                let dst_code = RegisterEncoding::get_16bit_register_code(dst)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", dst),
                        line: 0,
                    })?;
                let src_code = RegisterEncoding::get_16bit_register_code(src)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", src),
                        line: 0,
                    })?;
                // TEST reg, reg (16位) = 0x85
                let opcode = 0x85;
                let modrm = (0b11 << 6) | (dst_code << 3) | src_code;
                Ok(vec![opcode, modrm])
            }
            // TEST reg, imm
            (Operand::Register(reg), Operand::Immediate(imm)) => {
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                
                // 检查立即数范围
                if *imm < -32768 || *imm > 65535 {
                    return Err(AssemblerError::OperandError {
                        message: format!("立即数超出16位范围: {}", imm),
                        line: 0,
                    });
                }

                // TEST reg, imm (16位) = 0xF7 /0
                // reg字段=000表示TEST
                let opcode = 0xF7;
                let modrm = (0b11 << 6) | (0b000 << 3) | reg_code;
                let imm_u16 = *imm as u16;
                Ok(vec![opcode, modrm, imm_u16 as u8, (imm_u16 >> 8) as u8])
            }
            // TEST reg, mem
            (Operand::Register(reg), Operand::Memory(mem)) => {
                use super::memory::MemoryEncoder;
                
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;

                // TEST reg, mem (16位) = 0x85
                let opcode = 0x85;
                let (modrm, disp_bytes) = MemoryEncoder::encode_memory_operand(mem, reg_code)?;

                let mut result = vec![opcode, modrm];
                result.extend_from_slice(&disp_bytes);
                Ok(result)
            }
            // TEST mem, reg
            (Operand::Memory(mem), Operand::Register(reg)) => {
                use super::memory::MemoryEncoder;
                
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;

                // TEST mem, reg (16位) = 0x85
                // 注意：TEST mem, reg 和 TEST reg, mem 编码相同（测试操作）
                let opcode = 0x85;
                let (modrm, disp_bytes) = MemoryEncoder::encode_memory_operand(mem, reg_code)?;

                let mut result = vec![opcode, modrm];
                result.extend_from_slice(&disp_bytes);
                Ok(result)
            }
            _ => Err(AssemblerError::OperandError {
                message: "测试 暂不支持此操作数组合".to_string(),
                line: 0,
            }),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encode_and_reg_reg() {
        let encoder = LogicalEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Register("基址".to_string()),
        ];
        let result = encoder.encode("与", &operands).unwrap();
        // AND AX, BX = 21 C3 (dst=AX=0, src=BX=3, modrm=0b11000011=0xC3)
        assert_eq!(result, vec![0x21, 0xC3]);
    }

    #[test]
    fn test_encode_or_reg_reg() {
        let encoder = LogicalEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Register("基址".to_string()),
        ];
        let result = encoder.encode("或", &operands).unwrap();
        // OR AX, BX = 09 C3 (dst=AX=0, src=BX=3, modrm=0b11000011=0xC3)
        assert_eq!(result, vec![0x09, 0xC3]);
    }

    #[test]
    fn test_encode_xor_reg_reg() {
        let encoder = LogicalEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Register("基址".to_string()),
        ];
        let result = encoder.encode("异或", &operands).unwrap();
        // XOR AX, BX = 31 C3 (dst=AX=0, src=BX=3, modrm=0b11000011=0xC3)
        assert_eq!(result, vec![0x31, 0xC3]);
    }

    #[test]
    fn test_encode_not() {
        let encoder = LogicalEncoder;
        let operands = vec![Operand::Register("累加器".to_string())];
        let result = encoder.encode("非", &operands).unwrap();
        // NOT AX = F7 D0
        assert_eq!(result, vec![0xF7, 0xD0]);
    }

    #[test]
    fn test_encode_and_reg_imm8() {
        let encoder = LogicalEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Immediate(0x7F), // 127，在8位范围内
        ];
        let result = encoder.encode("与", &operands).unwrap();
        // AND AX, 0x7F (8位立即数) = 83 E0 7F
        // modrm = (0b11 << 6) | (0b100 << 3) | 0 = 0b11100000 = 0xE0
        assert_eq!(result, vec![0x83, 0xE0, 0x7F]);
    }

    #[test]
    fn test_encode_and_reg_imm16() {
        let encoder = LogicalEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Immediate(0xFF), // 255，超出8位范围，使用16位编码
        ];
        let result = encoder.encode("与", &operands).unwrap();
        // AND AX, 0xFF (16位立即数) = 81 E0 FF 00
        // modrm = (0b11 << 6) | (0b100 << 3) | 0 = 0b11100000 = 0xE0
        assert_eq!(result, vec![0x81, 0xE0, 0xFF, 0x00]);
    }

    #[test]
    fn test_encode_test_reg_reg() {
        let encoder = LogicalEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Register("基址".to_string()),
        ];
        let result = encoder.encode("测试", &operands).unwrap();
        // TEST AX, BX = 85 C3 (dst=AX=0, src=BX=3, modrm=0b11000011=0xC3)
        assert_eq!(result, vec![0x85, 0xC3]);
    }

    #[test]
    fn test_encode_test_reg_imm() {
        let encoder = LogicalEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Immediate(0x1234),
        ];
        let result = encoder.encode("测试", &operands).unwrap();
        // TEST AX, 0x1234 = F7 C0 34 12
        // modrm = (0b11 << 6) | (0b000 << 3) | 0 = 0b11000000 = 0xC0
        assert_eq!(result, vec![0xF7, 0xC0, 0x34, 0x12]);
    }

    #[test]
    fn test_encode_test_reg_mem() {
        use crate::ast::MemoryOperand;
        let encoder = LogicalEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Memory(MemoryOperand {
                base: Some("基址".to_string()),
                index: None,
                scale: None,
                displacement: None,
            }),
        ];
        let result = encoder.encode("测试", &operands).unwrap();
        // TEST AX, [BX] = 85 07
        assert_eq!(result, vec![0x85, 0x07]);
    }
}

